
#include <windows.h>
#include <mfapi.h>
#include <Mfidl.h>
#include <mfreadwrite.h>
#include <Mferror.h>

#include <shlwapi.h>

#include <stdio.h>

#include "audio_util.hpp"

void EnumAudioInDevices(void)
{
    printf("Enumerate all audio-in device:\n");
    
	// Create an attribute store to hold the search criteria.
	IMFAttributes *mf_config = NULL;
	HRESULT hr = MFCreateAttributes(&mf_config, 1);

	// Request video capture devices.
	if (SUCCEEDED(hr))
	{
		hr = mf_config->SetGUID(MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_AUDCAP_GUID);
	}

	// Enumerate the devices,
	IMFActivate **mf_devices = NULL;
	UINT32 mf_devices_count = 0;
	if (SUCCEEDED(hr))
	{
		hr = MFEnumDeviceSources(mf_config, &mf_devices, &mf_devices_count);
	}

#if 0
	// Create a media source for the first device in the list.
	if (SUCCEEDED(hr))
	{
		if (mf_devices_count > 0)
		{
			hr = mf_devices[0]->ActivateObject(IID_PPV_ARGS(ppSource));
		}
		else
		{
			hr = MF_E_NOT_FOUND;
		}
	}
#endif

    char dev_name[256];
	for (DWORD i = 0; i < mf_devices_count; i++)
	{
		WCHAR *name;
		mf_devices[i]->GetAllocatedString(MF_DEVSOURCE_ATTRIBUTE_FRIENDLY_NAME, &name, NULL);
        
        WideCharToMultiByte(CP_ACP, WC_NO_BEST_FIT_CHARS, name, -1, dev_name, 256, NULL, NULL);
        
		printf("  Friendly name: %s\n", dev_name);
		CoTaskMemFree(name);
	}

	for (DWORD i = 0; i < mf_devices_count; i++)
	{
		mf_devices[i]->Release();
	}
	CoTaskMemFree(mf_devices);
	SafeRelease(mf_config);

    return;
}

int GetAudioInDevice(IMFMediaSource **media_source, wchar_t *dev_name)
{
	IMFMediaSource *source = NULL;

	// Create an attribute store to hold the search criteria.
	IMFAttributes *mf_config = NULL;
	HRESULT hr = MFCreateAttributes(&mf_config, 1);

	// Request video capture devices.
	hr = mf_config->SetGUID(MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE, MF_DEVSOURCE_ATTRIBUTE_SOURCE_TYPE_AUDCAP_GUID);

	// Enumerate the devices,
	IMFActivate **mf_devices = NULL;
	UINT32 mf_devices_count = 0;
	hr = MFEnumDeviceSources(mf_config, &mf_devices, &mf_devices_count);
	SafeRelease(mf_config);
    
    if (mf_devices_count == 0)
    {
        return -1;
    }

    if (dev_name)
    {
        for (UINT32 i = 0; i < mf_devices_count; ++i)
        {
            WCHAR *name;
            mf_devices[i]->GetAllocatedString(MF_DEVSOURCE_ATTRIBUTE_FRIENDLY_NAME, &name, NULL);
            if (wcscmp(name, dev_name) == 0)
            {
                hr = mf_devices[i]->ActivateObject(IID_PPV_ARGS(media_source));
            }
            CoTaskMemFree(name);
            
            if (*media_source) { break; }
        }
    }
    else
    {
        hr = mf_devices[0]->ActivateObject(IID_PPV_ARGS(media_source));
    }

	for (DWORD i = 0; i < mf_devices_count; i++)
	{
		mf_devices[i]->Release();
	}
	CoTaskMemFree(mf_devices);

	return 0;
}

/************************************************************/

class AudioSampleSink : public IMFSourceReaderCallback
{
public:
	AudioSampleSink() : ref_count_(0), media_reader_(NULL) {}
	~AudioSampleSink() {}

	// IUnknown methods
	STDMETHODIMP QueryInterface(REFIID iid, void** ppv)
	{
		static const QITAB qit[] =
		{
			QITABENT(AudioSampleSink, IMFSourceReaderCallback),
			{ 0 },
		};
		return QISearch(this, qit, iid, ppv);
	}
	STDMETHODIMP_(ULONG) AddRef()
	{
		return InterlockedIncrement(&ref_count_);
	}
	STDMETHODIMP_(ULONG) Release()
	{
		ULONG uCount = InterlockedDecrement(&ref_count_);
		if (uCount == 0)
		{
			delete this;
		}
		return uCount;
	}

	HRESULT OnEvent(DWORD dwStreamIndex, IMFMediaEvent *pEvent)
	{
		return S_OK;
	}

	HRESULT OnFlush(DWORD dwStreamIndex)
	{
		return S_OK;
	}

	HRESULT OnReadSample(HRESULT hrStatus, DWORD dwStreamIndex, DWORD dwStreamFlags, LONGLONG llTimestamp, IMFSample *pSample)
	{
		printf("Audio Sample, ts: %I64d\n", llTimestamp);

		media_reader_->ReadSample(MF_SOURCE_READER_FIRST_AUDIO_STREAM, 0, NULL, NULL, NULL, NULL);

		return S_OK;
	}

	void SetSourceReader(IMFSourceReader *reader)
	{
		media_reader_ = reader;
	}

private:
	long ref_count_;
	IMFSourceReader *media_reader_;
};

int SetupAudioCaptureAsync(IMFMediaSource *source, IMFSourceReader **reader, DWORD *reader_stream_index, AudioSampleSink *sink)
{
	HRESULT hr;
	IMFSourceReader *media_reader;
	int media_reader_ready = 0;

	IMFAttributes *attr;
	MFCreateAttributes(&attr, 1);
	attr->SetUnknown(MF_SOURCE_READER_ASYNC_CALLBACK, sink);

	hr = MFCreateSourceReaderFromMediaSource(source, attr, &media_reader);

	SafeRelease(attr);

	DWORD desired_audio_stream = (DWORD)-1;

	/* enumerate all media stream type */
	DWORD stream_index = 0;
	do
	{
		IMFMediaType *media_type = NULL;
		hr = media_reader->GetCurrentMediaType(stream_index, &media_type);
		if (hr == S_OK)
		{
			GUID major_type;
			media_type->GetMajorType(&major_type);
			if (major_type == MFMediaType_Audio)
			{
				printf("Audio Stream, index: %u\n", stream_index);
				desired_audio_stream = stream_index;
				SafeRelease(media_type);
				break;
			}
			else if (major_type == MFMediaType_Video)
			{
				printf("Video Stream, index: %u\n", stream_index);
			}
			else
			{
				printf("Other Stream, index: %u\n", stream_index);
			}
			SafeRelease(media_type);
			stream_index++;
		}
	} while (hr == S_OK);

	if (desired_audio_stream == (DWORD)-1)
	{
		printf("No Desored Audio Stream was found\n");
		return -1;
	}

	DWORD media_type_index = 0;
	do
	{
		IMFMediaType *media_type = NULL;
		hr = media_reader->GetNativeMediaType(desired_audio_stream, media_type_index, &media_type);
		if (hr == S_OK)
		{
			media_type->SetGUID(MF_MT_SUBTYPE, MFAudioFormat_PCM);
			hr = media_reader->SetCurrentMediaType(desired_audio_stream, NULL, media_type);
			DumpAudioMediaType(media_type);
			SafeRelease(media_type);

			if (hr == S_OK)
			{
				break;
			}
			media_type_index++;
		}
	} while (hr == S_OK);

	if (hr != S_OK)
	{
		printf("Failed to get a desired media reader\n");
		return -2;
	}

	sink->SetSourceReader(media_reader);
	*reader_stream_index = desired_audio_stream;
	*reader = media_reader;

	return 0;
}

void AudioCaptureAsync(wchar_t *dev_name)
{
	IMFMediaSource *media_source = NULL;
	GetAudioInDevice(&media_source, dev_name);

	AudioSampleSink audio_sink;
	IMFSourceReader *media_reader = NULL;
	DWORD media_stream_index;
	SetupAudioCaptureAsync(media_source, &media_reader, &media_stream_index, &audio_sink);

	IMFPresentationDescriptor *desc = NULL;
	media_source->CreatePresentationDescriptor(&desc);


	PROPVARIANT var;
	PropVariantInit(&var);
	var.vt = VT_EMPTY;
	media_source->Start(desc, NULL, &var);

	media_reader->ReadSample(MF_SOURCE_READER_FIRST_AUDIO_STREAM, 0, NULL, NULL, NULL, NULL);

	getchar();

	media_source->Stop();
	media_source->Shutdown();
	SafeRelease(media_reader);
	SafeRelease(media_source);
}

/************************************************************/

int SetupAudioReaderFromSource(IMFMediaSource *source, IMFSourceReader **reader, DWORD *stream_index_out)
{
	HRESULT hr;
	IMFSourceReader *media_reader;

	hr = MFCreateSourceReaderFromMediaSource(source, NULL, &media_reader);
    
    EnumAllMediaTypes(media_reader);

    DWORD select_stream_index = (DWORD)-1;
    IMFMediaType *select_media_type = NULL;
    
    UINT32 max_sample_rate = 0;
    DWORD stream_index;
    for (stream_index = 0;; stream_index++)
    {
        DWORD type_index;
        for (type_index = 0;;type_index++)
        {
            IMFMediaType *media_type = NULL;
            hr = media_reader->GetNativeMediaType(stream_index, type_index, &media_type);
            if (hr != S_OK) { break; }

            GUID maj_type;
            media_type->GetMajorType(&maj_type);
            if (maj_type == MFMediaType_Audio)
            {
                GUID sub_type;
                media_type->GetGUID(MF_MT_SUBTYPE, &sub_type);
                if (sub_type == MFAudioFormat_Float || sub_type == MFAudioFormat_PCM)
                {
                    UINT32 sample_rate = 0;
                    media_type->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &sample_rate);
                    if (sample_rate > max_sample_rate)
                    {
                        media_type->AddRef();
                        SafeRelease(select_media_type);
                        select_stream_index = stream_index;
                        select_media_type = media_type;
                    }
                }
            }

            media_type->Release();
        }

        if (hr == MF_E_INVALIDSTREAMNUMBER) { break; }
    }
    
    if (select_media_type == NULL)
    {
        printf("no pcm media stream available\n");
        media_reader->Release();
        return -1;
    }

    UINT32 channels = 0;
    select_media_type->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &channels);
    UINT32 samples_rate = 0;
    select_media_type->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &samples_rate);
    UINT32 sample_bits = 16;
    select_media_type->SetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, sample_bits);
    UINT32 block_align = channels * samples_rate * (sample_bits >> 3);
    select_media_type->SetUINT32(MF_MT_AUDIO_BLOCK_ALIGNMENT, block_align);
    UINT32 avg_bytes_per_sec = block_align;
    select_media_type->SetUINT32(MF_MT_AUDIO_AVG_BYTES_PER_SECOND, avg_bytes_per_sec);
    hr = media_reader->SetCurrentMediaType(select_stream_index, NULL, select_media_type);

	*reader = media_reader;
    *stream_index_out = select_stream_index;

	return 0;
}

int cap_run = 0;
BOOL WINAPI OnConsoleCtrl(DWORD dwCtrlType) 
{
	if (dwCtrlType == CTRL_C_EVENT)
	{
		cap_run = 0;
		return TRUE;
	}

	return FALSE;
}

void AudioCaptureSync(wchar_t *dev_name)
{
	IMFMediaSource *media_source = NULL;
    GetAudioInDevice(&media_source, dev_name);

	IMFSourceReader *media_reader = NULL;
    DWORD stream_index = (DWORD)-1;
	SetupAudioReaderFromSource(media_source, &media_reader, &stream_index);

	IMFPresentationDescriptor *desc = NULL;
	media_source->CreatePresentationDescriptor(&desc);

	PROPVARIANT var;
	PropVariantInit(&var);
	var.vt = VT_EMPTY;
	media_source->Start(desc, NULL, &var);

	cap_run = 1;
	SetConsoleCtrlHandler(OnConsoleCtrl, TRUE);
	while (cap_run)
	{
		DWORD stream_idx_got, flags;
		LONGLONG llTimeStamp;
		IMFSample *sample;
		HRESULT hr = media_reader->ReadSample(stream_index, 0, &stream_idx_got, &flags, &llTimeStamp, &sample);
		if (hr == S_OK)
		{
			printf("Audio Sample, stream_index: %u, stream_flag: 0x%X, pts: %I64dms", stream_idx_got, flags, (llTimeStamp / 10000));
			if (sample) 
			{
				DWORD data_len = 0;
				sample->GetTotalLength(&data_len);
				DWORD buf_count = 0;
				sample->GetBufferCount(&buf_count);
				LONGLONG sample_duration;
				sample->GetSampleDuration(&sample_duration);

				sample->Release();

				printf(", data_len: %lu, buf_count: %lu, duration: %I64dms\n", data_len, buf_count, (sample_duration / 10000));
			}
			else
			{
				printf("\n");
			}
		}
	}

	media_source->Stop();
	media_source->Shutdown();
	SafeRelease(media_reader);
	SafeRelease(media_source);
}

/************************************************************/

void audio_cap_routine(int argc, wchar_t *argv[])
{
    CoInitializeEx(0, COINIT_MULTITHREADED);
    MFStartup(MF_VERSION);

	EnumAudioInDevices();

	AudioCaptureSync(L"麦克风 (USB Microphone)");
    //AudioCaptureSync(L"麦克风阵列 (Realtek(R) Audio)");

    MFShutdown();
    CoUninitialize();
    
    return;
}

int wmain(int argc, wchar_t *argv[])
{
    audio_cap_routine(argc, argv);

    return 0;
}
